import {
  createProviderToolFactoryWithOutputSchema,
  lazySchema,
  zodSchema,
} from '@ai-sdk/provider-utils';
import { JSONValue } from '@ai-sdk/provider';
import { z } from 'zod/v4';

const jsonValueSchema: z.ZodType<JSONValue> = z.lazy(() =>
  z.union([
    z.string(),
    z.number(),
    z.boolean(),
    z.null(),
    z.array(jsonValueSchema),
    z.record(z.string(), jsonValueSchema),
  ]),
);

export const mcpArgsSchema = lazySchema(() =>
  zodSchema(
    z
      .object({
        serverLabel: z.string(),
        allowedTools: z
          .union([
            z.array(z.string()),
            z.object({
              readOnly: z.boolean().optional(),
              toolNames: z.array(z.string()).optional(),
            }),
          ])
          .optional(),
        authorization: z.string().optional(),
        connectorId: z.string().optional(),
        headers: z.record(z.string(), z.string()).optional(),

        // TODO: Integrate this MCP tool approval with our SDK's existing tool approval architecture
        // requireApproval: z
        //   .union([
        //     z.enum(['always', 'never']),
        //     z.object({
        //       readOnly: z.boolean().optional(),
        //       toolNames: z.array(z.string()).optional(),
        //     }),
        //   ])
        //   .optional(),
        serverDescription: z.string().optional(),
        serverUrl: z.string().optional(),
      })
      .refine(
        v => v.serverUrl != null || v.connectorId != null,
        'One of serverUrl or connectorId must be provided.',
      ),
  ),
);

const mcpInputSchema = lazySchema(() => zodSchema(z.object({})));

export const mcpOutputSchema = lazySchema(() =>
  zodSchema(
    z.discriminatedUnion('type', [
      z.object({
        type: z.literal('call'),
        serverLabel: z.string(),
        name: z.string(),
        arguments: z.string(),
        output: z.string().nullable().optional(),
        error: z.union([z.string(), jsonValueSchema]).optional(),
      }),
      z.object({
        type: z.literal('listTools'),
        serverLabel: z.string(),
        tools: z.array(
          z.object({
            name: z.string(),
            description: z.string().optional(),
            inputSchema: jsonValueSchema,
            annotations: z.record(z.string(), jsonValueSchema).optional(),
          }),
        ),
        error: z.union([z.string(), jsonValueSchema]).optional(),
      }),
      z.object({
        type: z.literal('approvalRequest'),
        serverLabel: z.string(),
        name: z.string(),
        arguments: z.string(),
        approvalRequestId: z.string(),
      }),
    ]),
  ),
);

type McpArgs = {
  /** A label for this MCP server, used to identify it in tool calls. */
  serverLabel: string;
  /** List of allowed tool names or a filter object. */
  allowedTools?:
    | string[]
    | {
        readOnly?: boolean;
        toolNames?: string[];
      };
  /** OAuth access token usable with the remote MCP server or connector. */
  authorization?: string;
  /** Identifier for a service connector. */
  connectorId?: string;
  /** Optional HTTP headers to send to the MCP server. */
  headers?: Record<string, string>;
  // /** Which tools require approval: 'always' | 'never' | filter object. */
  // requireApproval?:
  //   | 'always'
  //   | 'never'
  //   | {
  //       readOnly?: boolean;
  //       toolNames?: string[];
  //     };
  /** Optional description of the MCP server. */
  serverDescription?: string;
  /** URL for the MCP server. One of serverUrl or connectorId must be provided. */
  serverUrl?: string;
};

export const mcpToolFactory = createProviderToolFactoryWithOutputSchema<
  {},
  | {
      type: 'call';
      serverLabel: string;
      name: string;
      arguments: string;
      output?: string | null;
      error?: JSONValue;
    }
  | {
      type: 'listTools';
      serverLabel: string;
      tools: Array<{
        name: string;
        description?: string;
        inputSchema: unknown;
        annotations?: unknown;
      }>;
      error?: JSONValue;
    }
  | {
      type: 'approvalRequest';
      serverLabel: string;
      name: string;
      arguments: string;
      approvalRequestId: string;
    },
  McpArgs
>({
  id: 'openai.mcp',
  inputSchema: mcpInputSchema,
  outputSchema: mcpOutputSchema,
});

export const mcp = (args: McpArgs) => mcpToolFactory(args);
